import { ChatAgent, ChatAgentResponse } from "../agent/chatAgent";
import { Language, RoleMessageType } from "../enum";
import { ReceivingChatMessage } from "../message/chatMessage";
import { getPromptFromKey, promptTemplateDict } from "../prompt";

/**
 * 表示两个Agent之间进行对话
 * @param initiator 对话发起者
 * @param responder 对话接受者
 * @param topic 对话主题
 * @returns 对话双方 [initiator, responder]， 如果对话失败则返回undefined
 *
 * @description
 *
 * 1. 首先判断两个Agent是否可以进行对话，如果不能则返回undefined
 * 2. 如果有对话主题，则先将主题发送给对话发起者，对话发起者再发送消息给对话接受者
 * 3. 如果没有对话主题，则直接发送对话发起者的最新消息给对话接受者
 */
export async function talk(
    initiator: ChatAgent,
    responder: ChatAgent,
    topic?: ReceivingChatMessage,
    useMemory: boolean = true,
    memoryDepth: number = 1
): Promise<{ initiatorRes?: ChatAgentResponse; responderRes?: ChatAgentResponse }> {
    if (!canTalk(initiator) || !canTalk(responder)) {
        return {};
    }

    const response: {
        initiatorRes?: ChatAgentResponse;
        responderRes?: ChatAgentResponse;
    } = {};

    if (topic) {
        const content = getTalkContent({
            otherRoleName: responder.roleName,
            language: initiator.language,
            topic: topic.content,
            agent: initiator,
            useMemory,
            memoryDepth
        });
        const startMessate = new ReceivingChatMessage({
            ...topic.getConfig(),
            roleMessageType: RoleMessageType.AgentReceiving,
            content
        });
        // 发送对话主题给对话发起者
        response.initiatorRes = await initiator.step(startMessate);
        await talk(initiator, responder, undefined);
    } else {
        const latestInitiatorMessage = initiator.getLatestMessage();
        const content = getTalkContent({
            otherRoleName: initiator.roleName,
            language: responder.language,
            content: latestInitiatorMessage.content,
            agent: responder,
            useMemory,
            memoryDepth
        });
        const startMessate = new ReceivingChatMessage({
            ...latestInitiatorMessage.getConfig(),
            roleMessageType: RoleMessageType.AgentReceiving,
            content
        });
        // 发送对话发起者的最新消息给对话接受者
        response.responderRes = await responder.step(startMessate);
    }

    if (initiator.terminated || responder.terminated) {
        return {};
    }

    return response;
}

function getTalkContent(prompt: {
    otherRoleName: string;
    language: Language;
    useMemory: boolean;
    memoryDepth: number;
    agent: ChatAgent;
    topic?: string;
    content?: string;
}) {
    const talkHead = getPromptFromKey(
        promptTemplateDict.baseTemplate.talkHead,
        prompt.language
    ).format({
        role: prompt.otherRoleName
    });
    const talkTopic = prompt.topic
        ? getPromptFromKey(promptTemplateDict.baseTemplate.talkTopic, prompt.language).format({
              topic: prompt.topic
          })
        : "";
    const talkContent = prompt.content
        ? getPromptFromKey(promptTemplateDict.baseTemplate.talkContent, prompt.language).format({
              content: prompt.content
          })
        : "";
    const content = talkHead + "\n" + talkTopic + "\n" + talkContent;
    const agent = prompt.agent;
    const talkMemory = prompt.useMemory
        ? getPromptFromKey(promptTemplateDict.baseTemplate.roleMemory, prompt.language).format({
              memory: agent
                  .useMemory(
                      new ReceivingChatMessage({
                          roleName: agent.roleName,
                          roleType: agent.roleType,
                          content
                      }),
                      prompt.memoryDepth
                  )
                  .reduce((acc, cur) => acc + cur.content, "")
          })
        : "";
    return talkHead + "\n" + talkTopic + "\n" + talkContent + "\n" + talkMemory;
}

function canTalk(agent: ChatAgent) {
    if (agent.terminated) {
        return false;
    }
    if (agent.getLatestMessage().roleMessageType === RoleMessageType.AgentReceiving) {
        return false;
    }
    return true;
}
